import pandas as pd
import numpy as np
from config import *
import torch

def load_data(F_X, F_y):
    """
    加载训练集数据，并制作数据加载器
    """
    # 获取所有的特征名称， 并转为列表
    # featureNames = pd.read_csv(F_features, sep=' ', header=None, skipinitialspace=True)[1].tolist()
    # 获取所有的数据标签， 并转为列表
    y_train_list = pd.read_csv(F_y, sep=' ', header=None, skipinitialspace=True)[0]
    y_train_list = [int(i) - 1 for i in y_train_list]  # 减1，因为标签从0开始
    # 获取训练集数据
    df_X_train = pd.read_csv(F_X, sep=' ', header=None, skipinitialspace=True)

    # 将数据转换为数组
    features = np.array(df_X_train)
    labels = np.array(y_train_list)

    # 将数据转为PyTorch张量
    features = torch.tensor(features, dtype=torch.float32)
    labels = torch.tensor(labels, dtype=torch.long)

    # 创建数据集，并制作数据加载器
    dataset = torch.utils.data.TensorDataset(features, labels)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)

    return dataloader
